package com.xzy.mybatis.cipher.spring;

import com.xzy.mybatis.cipher.annotation.CryptoColumn;
import com.xzy.mybatis.cipher.annotation.CryptoTable;
import com.xzy.mybatis.cipher.annotation.CustomCryptoStrategy;
import com.xzy.mybatis.cipher.annotation.EnableMybatisCipher;
import com.xzy.mybatis.cipher.support.enums.ParseStrategyEnum;
import com.xzy.mybatis.cipher.support.exceptions.MyBatisCipherException;
import com.xzy.mybatis.cipher.support.exceptions.NotSpecifiedCryptoStrategyException;
import com.xzy.mybatis.cipher.support.exceptions.UnsupportedCryptoStrategyException;
import com.xzy.mybatis.cipher.support.metadata.CryptoColumnMetadata;
import com.xzy.mybatis.cipher.support.metadata.CryptoTableMetadata;
import com.xzy.mybatis.cipher.support.properties.*;
import com.xzy.mybatis.cipher.support.registry.CryptoMetadataRegistry;
import com.xzy.mybatis.cipher.support.registry.DataCryptoStrategyRegistry;
import com.xzy.mybatis.cipher.support.registry.SqlParseStrategyRegistry;
import com.xzy.mybatis.cipher.support.strategy.crypto.CryptoStrategy;
import com.xzy.mybatis.cipher.support.strategy.crypto.CryptoStrategyFactory;
import com.xzy.mybatis.cipher.support.strategy.parse.DeleteParseStrategy;
import com.xzy.mybatis.cipher.support.strategy.parse.InsertParseStrategy;
import com.xzy.mybatis.cipher.support.strategy.parse.SelectParseStrategy;
import com.xzy.mybatis.cipher.support.strategy.parse.UpdateParseStrategy;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.support.BeanDefinitionBuilder;
import org.springframework.beans.factory.support.BeanDefinitionRegistry;
import org.springframework.boot.context.properties.bind.Bindable;
import org.springframework.boot.context.properties.bind.Binder;
import org.springframework.context.EnvironmentAware;
import org.springframework.context.annotation.ImportBeanDefinitionRegistrar;
import org.springframework.core.annotation.AnnotationAttributes;
import org.springframework.core.env.Environment;
import org.springframework.core.io.Resource;
import org.springframework.core.io.support.PathMatchingResourcePatternResolver;
import org.springframework.core.io.support.ResourcePatternResolver;
import org.springframework.core.type.AnnotationMetadata;
import org.springframework.core.type.classreading.CachingMetadataReaderFactory;
import org.springframework.core.type.classreading.MetadataReader;
import org.springframework.core.type.classreading.MetadataReaderFactory;
import org.springframework.util.ClassUtils;

import java.io.IOException;
import java.lang.reflect.Field;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Objects;

/**
 * @author xzy.xiao
 * @since 2025/1/14  18:48
 */
@Slf4j
public class MybatisCipherRegistrar implements ImportBeanDefinitionRegistrar, EnvironmentAware {

    private Environment environment;

    @Override
    public void setEnvironment(Environment environment) {
        this.environment = environment;
    }

    @Override
    public void registerBeanDefinitions(AnnotationMetadata importingClassMetadata, BeanDefinitionRegistry registry) {
        log.debug("1. Init context ...");
        MybatisCipherContext mybatisCipherContext = initContext(importingClassMetadata);

        log.debug("2. Register plugin ...");
        registerPlugin(registry, mybatisCipherContext);
    }

    private MybatisCipherContext initContext(AnnotationMetadata importingClassMetadata) {
        Map<String, Object> annotationAttributes = importingClassMetadata.getAnnotationAttributes(EnableMybatisCipher.class.getName());
        if (null == annotationAttributes) {
            log.error("The '@EnableMybatisCipher' annotation is missing !");
            throw new MyBatisCipherException("The '@EnableMybatisCipher' annotation is missing !");
        }

        String[] scanPackages = (String[]) annotationAttributes.get("scanPackages");
        AnnotationAttributes[] customCryptoStrategies = (AnnotationAttributes[]) annotationAttributes.get("customCryptoStrategies");

        log.debug("1.1 Init MybatisCipherProperties ...");
        MybatisCipherProperties mybatisCipherProperties = initMybatisCipherProperties(environment);

        log.debug("1.2 Init CryptoMetadataRegistry ...");
        CryptoMetadataRegistry cryptoMetadataRegistry = initCryptoMetadataRegistry(scanPackages);

        log.debug("1.3 Init DataCryptoStrategyRegistry ...");
        DataCryptoStrategyRegistry dataCryptoStrategyRegistry = initDataCryptoStrategyRegistry(customCryptoStrategies, mybatisCipherProperties, cryptoMetadataRegistry);

        log.debug("1.4 Init SqlParseStrategyRegistry ...");
        SqlParseStrategyRegistry sqlParseStrategyRegistry = initSqlParseStrategyRegistry(cryptoMetadataRegistry, dataCryptoStrategyRegistry);

        return new MybatisCipherContext(mybatisCipherProperties, cryptoMetadataRegistry, dataCryptoStrategyRegistry, sqlParseStrategyRegistry);
    }

    private MybatisCipherProperties initMybatisCipherProperties(Environment environment) {
        // Init
        MybatisCipherProperties mybatisCipherProperties = Binder.get(environment)
                .bind(MybatisCipherProperties.PREFIX, Bindable.of(MybatisCipherProperties.class))
                .orElse(null);

        // Check
        if (null == mybatisCipherProperties) {
            log.error("Stop init mybatis cipher properties, because the 'mybatis-cipher' properties is missing!");
            throw new MyBatisCipherException("Stop init mybatis cipher properties, because the 'mybatis-cipher' properties is missing!");
        }

        Map<String, DatasourceProperty> datasourcePropertyMap = mybatisCipherProperties.getDatasource();
        if (null == datasourcePropertyMap || datasourcePropertyMap.isEmpty()) {
            log.error("Stop init mybatis cipher properties, because the datasource properties is empty!");
            throw new MyBatisCipherException("Stop init mybatis cipher properties, because the datasource properties is empty!");
        }

        for (String datasourceName : datasourcePropertyMap.keySet()) {
            DatasourceProperty datasourceProperty = datasourcePropertyMap.get(datasourceName);
            PrimaryProperty primaryProperty = datasourceProperty.getPrimary();
            Map<String, TableProperty> tablePropertyMap = datasourceProperty.getTable();

            if (null == primaryProperty && (null == tablePropertyMap || tablePropertyMap.isEmpty())) {
                log.error("Stop init mybatis cipher properties, because the datasource properties is empty!");
                throw new MyBatisCipherException("Stop init mybatis cipher properties, because the datasource properties is empty!");
            }
        }

        return mybatisCipherProperties;
    }

    private CryptoMetadataRegistry initCryptoMetadataRegistry(String[] scanPackages) {
        CryptoMetadataRegistry cryptoMetadataRegistry = new CryptoMetadataRegistry();

        for (String scanPackage : scanPackages) {
            try {
                searchCryptoMetadata(cryptoMetadataRegistry, scanPackage);
            } catch (Exception e) {
                log.error("Searching for crypto metadata failed for package: {}", scanPackage);
                throw new MyBatisCipherException(e);
            }
        }

        return cryptoMetadataRegistry;
    }

    private void searchCryptoMetadata(CryptoMetadataRegistry cryptoMetadataRegistry, String scanPackage) throws Exception {
        log.debug("Searching for crypto metadata for package: {}", scanPackage);

        // Get the class file in the package
        ResourcePatternResolver resourcePatternResolver = new PathMatchingResourcePatternResolver();
        MetadataReaderFactory metadataReaderFactory = new CachingMetadataReaderFactory(resourcePatternResolver);

        String pattern = ResourcePatternResolver.CLASSPATH_ALL_URL_PREFIX + ClassUtils.convertClassNameToResourcePath(scanPackage) + "/**/*.class";
        Resource[] resources = resourcePatternResolver.getResources(pattern);

        // Searching for crypto metadata associated with each class
        for (Resource resource : resources) {
            CryptoTableMetadata cryptoMetadata = searchCryptoMetadata(resource, metadataReaderFactory);
            if (null != cryptoMetadata) {
                log.debug("Register crypto metadata：{}", cryptoMetadata);
                cryptoMetadataRegistry.registerTableMetadata(cryptoMetadata);
            }
        }
    }

    private CryptoTableMetadata searchCryptoMetadata(Resource resource, MetadataReaderFactory metadataReaderFactory) throws IOException, ClassNotFoundException {
        if (!resource.isReadable()) {
            return null;
        }

        MetadataReader metadataReader = metadataReaderFactory.getMetadataReader(resource);
        String className = metadataReader.getClassMetadata().getClassName();
        Class<?> clazz = Class.forName(className);
        log.debug("Searching crypto metadata on class: {}", className);

        // Table metadata
        CryptoTable cryptoTable = clazz.getAnnotation(CryptoTable.class);
        if (Objects.isNull(cryptoTable)) {
            return null;
        }
        CryptoTableMetadata tableMetadata = new CryptoTableMetadata();

        // Column metadata
        List<CryptoColumnMetadata> columnMetadataList = new ArrayList<>();
        for (Field field : clazz.getDeclaredFields()) {

            CryptoColumn cryptoColumn = field.getAnnotation(CryptoColumn.class);
            if (null == cryptoColumn) {
                continue;
            }

            if (!String.class.equals(field.getType())) {
                log.error("'@CryptoColumn' is above non-string. class:{}, field: {}", clazz, field.getName());
                throw new MyBatisCipherException("'@CryptoColumn' is above non-string field !");
            }

            columnMetadataList.add(new CryptoColumnMetadata()
                    .setFieldName(field.getName())
                    .setColumnName(cryptoColumn.columnName())
                    .setTableMetadata(tableMetadata)
            );
        }

        return tableMetadata
                .setClazz(clazz)
                .setClassName(className)
                .setTableName(cryptoTable.tableName())
                .setDatasourceName(cryptoTable.datasourceName())
                .setColumnMetadataList(columnMetadataList);
    }

    private DataCryptoStrategyRegistry initDataCryptoStrategyRegistry(AnnotationAttributes[] customCryptoStrategies, MybatisCipherProperties mybatisCipherProperties, CryptoMetadataRegistry cryptoMetadataRegistry) {
        DataCryptoStrategyRegistry dataCryptoStrategyRegistry = new DataCryptoStrategyRegistry(cryptoMetadataRegistry);

        // Register custom crypto strategy
        log.debug("Register custom crypto strategy ...");
        registerCustomCryptoStrategy(customCryptoStrategies);

        Map<String, DatasourceProperty> datasourceName2datasourcePropertyMap = mybatisCipherProperties.getDatasource();
        Map<String, Map<String, Map<String, CryptoColumnMetadata>>> datasourceName2cryptoColumnMetadataMap = cryptoMetadataRegistry.getDatasource2table2columnMetadataMap();
        for (String datasourceName : datasourceName2cryptoColumnMetadataMap.keySet()) {

            Map<String, Map<String, CryptoColumnMetadata>> tableName2cryptoColumnMetadataMap = datasourceName2cryptoColumnMetadataMap.get(datasourceName);
            DatasourceProperty datasourceProperty = datasourceName2datasourcePropertyMap.getOrDefault(datasourceName, null);
            if (null == datasourceProperty) {
                log.error("Stop init data crypto strategy registry, because data source property is empty! datasourceName:{}", datasourceName);
                throw new MyBatisCipherException("Stop init data crypto strategy registry, because data source property is empty! datasourceName:" + datasourceName);
            }

            // Register datasource primary data crypto strategy
            log.debug("Register datasource primary data crypto strategy ...");
            PrimaryProperty primaryProperty = datasourceProperty.getPrimary();
            registerPrimaryCryptoStrategy(dataCryptoStrategyRegistry, datasourceName, primaryProperty);

            // Register column data crypto strategy
            log.debug("Register column data crypto strategy ...");
            Map<String, TableProperty> tableName2tablePropertyMap = datasourceProperty.getTable();
            registerColumnCryptoStrategy(dataCryptoStrategyRegistry, datasourceName, tableName2cryptoColumnMetadataMap, tableName2tablePropertyMap);

        }

        return dataCryptoStrategyRegistry;
    }

    private static void registerCustomCryptoStrategy(AnnotationAttributes[] customCryptoStrategies) {
        for (AnnotationAttributes customCryptoStrategy : customCryptoStrategies) {
            String customCryptoStrategyName = (String) customCryptoStrategy.get("name");
            Class<? extends CryptoStrategy> customCryptoStrategyClass = (Class<? extends CryptoStrategy>) customCryptoStrategy.get("clazz");

            CryptoStrategyFactory.registerStrategy(customCryptoStrategyName, customCryptoStrategyClass);
        }
    }

    private void registerPrimaryCryptoStrategy(DataCryptoStrategyRegistry dataCryptoStrategyRegistry, String datasourceName, PrimaryProperty primaryProperty) {
        if (null == primaryProperty) {
            log.warn("Stop register datasource primary data crypto strategy for [{}], because the datasource primary properties is empty.", datasourceName);
            return;
        }

        String primaryCryptoStrategyName = primaryProperty.getStrategy();
        if (null == primaryCryptoStrategyName) {
            log.warn("Stop register datasource primary data crypto strategy for [{}], because the datasource primary strategy is empty.", datasourceName);
            return;
        }

        String primaryCryptoStrategyKey = primaryProperty.getKey();
        if (null == primaryCryptoStrategyKey) {
            log.warn("Stop register datasource primary data crypto strategy for [{}], because the datasource primary strategy key is empty.", datasourceName);
            return;
        }

        try {
            CryptoStrategy strategy = CryptoStrategyFactory.getStrategy(primaryCryptoStrategyName, primaryCryptoStrategyKey);
            dataCryptoStrategyRegistry.registerPrimaryStrategy(datasourceName, strategy);
        } catch (Exception e) {
            log.error("Stop register datasource primary data crypto strategy, because the data crypto strategy is not support! datasourceName:{}, cryptoStrategyName:{}", datasourceName, primaryCryptoStrategyName);
            throw new UnsupportedCryptoStrategyException("Unsupported crypto strategy: " + primaryCryptoStrategyName);
        }
    }

    private void registerColumnCryptoStrategy(DataCryptoStrategyRegistry dataCryptoStrategyRegistry, String datasourceName, Map<String, Map<String, CryptoColumnMetadata>> tableName2cryptoColumnMetadataMap, Map<String, TableProperty> tableName2tablePropertyMap) {
        for (String tableName : tableName2cryptoColumnMetadataMap.keySet()) {

            Map<String, CryptoColumnMetadata> columnName2cryptoColumnMetadataMap = tableName2cryptoColumnMetadataMap.get(tableName);
            TableProperty tableProperty = null != tableName2tablePropertyMap ? tableName2tablePropertyMap.getOrDefault(tableName, null) : null;

            for (String columnName : columnName2cryptoColumnMetadataMap.keySet()) {

                if (shouldUsePrimaryCryptoStrategy(tableProperty, columnName)) {
                    // Use datasource primary data crypto strategy
                    CryptoStrategy primaryStrategy = dataCryptoStrategyRegistry.getPrimaryStrategy(datasourceName);
                    if (null == primaryStrategy) {
                        log.error("Stop register table data crypto strategy, because the data crypto strategy is empty. datasourceName:{}, tableName:{}, columnName:{}", datasourceName, tableName, columnName);
                        throw new NotSpecifiedCryptoStrategyException("No specified crypto strategy: " + datasourceName + "." + tableName + "." + columnName);
                    } else {
                        dataCryptoStrategyRegistry.registerStrategy(datasourceName, tableName, columnName, primaryStrategy);
                    }
                } else {
                    // Use column data crypto strategy
                    ColumnProperty columnProperty = tableProperty.getColumn().get(columnName);
                    String columnCryptoStrategyName = columnProperty.getStrategy();
                    String columnCryptoStrategyKey = columnProperty.getKey();

                    try {
                        CryptoStrategy strategy = CryptoStrategyFactory.getStrategy(columnCryptoStrategyName, columnCryptoStrategyKey);
                        dataCryptoStrategyRegistry.registerStrategy(datasourceName, tableName, columnName, strategy);
                    } catch (Exception e) {
                        log.error("Stop register table data crypto strategy, because the data crypto strategy is not support! datasourceName:{}, tableName:{}, columnName:{}, cryptoStrategyName:{}", datasourceName, tableName, columnName, columnCryptoStrategyName);
                        throw new UnsupportedCryptoStrategyException("Unsupported crypto strategy: " + columnCryptoStrategyName);
                    }
                }

            }

        }
    }

    private boolean shouldUsePrimaryCryptoStrategy(TableProperty tableProperty, String columnName) {
        if (null == tableProperty) {
            return true;
        }

        Map<String, ColumnProperty> columnName2columnPropertyMap = tableProperty.getColumn();
        if (null == columnName2columnPropertyMap) {
            return true;
        }

        if (!columnName2columnPropertyMap.containsKey(columnName)) {
            return true;
        }

        ColumnProperty columnProperty = columnName2columnPropertyMap.get(columnName);
        return null == columnProperty.getStrategy() && null == columnProperty.getKey();
    }

    private SqlParseStrategyRegistry initSqlParseStrategyRegistry(CryptoMetadataRegistry cryptoMetadataRegistry, DataCryptoStrategyRegistry dataCryptoStrategyRegistry) {
        SqlParseStrategyRegistry sqlParseStrategyRegistry = new SqlParseStrategyRegistry();

        // Register system strategy
        log.debug("Register system sql parse strategy ...");
        sqlParseStrategyRegistry.registerStrategy(ParseStrategyEnum.INSERT.getCode(), new InsertParseStrategy(cryptoMetadataRegistry, dataCryptoStrategyRegistry));
        sqlParseStrategyRegistry.registerStrategy(ParseStrategyEnum.DELETE.getCode(), new DeleteParseStrategy(cryptoMetadataRegistry, dataCryptoStrategyRegistry));
        sqlParseStrategyRegistry.registerStrategy(ParseStrategyEnum.UPDATE.getCode(), new UpdateParseStrategy(cryptoMetadataRegistry, dataCryptoStrategyRegistry));
        sqlParseStrategyRegistry.registerStrategy(ParseStrategyEnum.SELECT.getCode(), new SelectParseStrategy(cryptoMetadataRegistry, dataCryptoStrategyRegistry));

        return sqlParseStrategyRegistry;
    }

    private void registerPlugin(BeanDefinitionRegistry registry, MybatisCipherContext mybatisCipherContext) {
        BeanDefinitionBuilder builder = BeanDefinitionBuilder.genericBeanDefinition(MybatisCipherInitializingBean.class);
        builder.addPropertyValue("mybatisCipherContext", mybatisCipherContext);
        registry.registerBeanDefinition("interceptorInitializingBean", builder.getBeanDefinition());
    }
}
